'''
This is a numpy representation of SAM format.
'''
import sys
import numpy as np
import h5py

# define the data type.

		
ham_50_dt = np.dtype([\
		('QNAME', 'S255'),\
		('FLAG', np.int),\
		('RNAME', 'S255'),\
		('POS', np.int),\
		('MAPQ', np.int),\
		('CIGAR', 'S50'),\
		('MRNM', 'S255'),\
		('MPOS', np.int),\
		('ISIZE', np.int),\
		('SEQ', 'S50'),\
		('QUAL', 'S50'),\
		('OPTIONAL','S50'),\
		])

ham_var_dt = np.dtype([\
		('QNAME', h5py.new_vlen(str)),\
		('FLAG', np.int),\
		('RNAME', h5py.new_vlen(str)),\
		('POS', np.int),\
		('MAPQ', np.int),\
		('CIGAR', h5py.new_vlen(str)),\
		('MRNM', h5py.new_vlen(str)),\
		('MPOS', np.int),\
		('ISIZE', np.int),\
		('SEQ', h5py.new_vlen(str)),\
		('QUAL', h5py.new_vlen(str)),\
		('OPTIONAL',h5py.new_vlen(str)),\
		])
		
ham_ant_dt = np.dtype([\
		('TYPE', np.int),\
		('MIDX', np.int),\
		])
		
initial_size = 10000000
default_chunk_size = 1000000
ham_dt = ham_50_dt

# create a numpy array to fill.
def create_array(sz):
	return np.zeros((sz,1), dtype=ham_dt)

# create a numpy array to fill.
def create_annotate_array(sz):
	return np.zeros((sz,1), dtype=ham_ant_dt)

class HamFile(object):
	
	# initializes h5 file and data set.
	def __init__(self, ham_file, init_size=False, chunk_size=False, cprs=False, read=0):
		# determine how to open.
		self.read = read
		if read == 0:
			self.open_write(ham_file, init_size, chunk_size, cprs)
		elif read == 1:
			self.open_read(ham_file)			
		else:
			self.open_both(ham_file)
	
	def open_both(self, ham_file):
		# open file for read and write.
		self.h5_in = h5py.File(ham_file, 'r+')
		
		# open data pointer.
		self.data = self.h5_in['nums']	
		
	def open_read(self, ham_file):
		# open file for read and write.
		self.h5_in = h5py.File(ham_file, 'r')
		
		# open data pointer.
		self.data = self.h5_in['nums']
			

	def open_write(self, ham_file, init_size, chunk_size, cprs):
		# See if initial size, or any of the chunk sizes were set.
		if init_size != False:
			self.initial_size = init_size
		else:
			self.initial_size = initial_size
			
		if chunk_size != False:
			self.chunk_size = chunk_size
		else:
			self.chunk_size = default_chunk_size
		
		# No compression if data type is vlen.
		if cprs == True and ham_dt == ham_var_dt:
			cprs = False
		
		# Open h5 file.
		self.h5_out = h5py.File(ham_file, 'w')
		
		# set data type.
		self.ham_dt = ham_dt
		
		# Create h5py object.
		if cprs == False:
			# without compression.
			self.data = self.h5_out.create_dataset('nums', (initial_size,1), self.ham_dt, maxshape=(None,1))
		else:
			# with compression.
			self.data = self.h5_out.create_dataset('nums', (initial_size,1), self.ham_dt, maxshape=(None,1), compression='gzip', chunks=(self.chunk_size,1))

		# Setup attributes.
		self.active_idx = 0
		self.max_idx = initial_size

	# closes h5 file.
	def close(self):
		if self.read == 0:
			# shrink if necessary.
			if self.max_idx > self.active_idx:
				# Check if we wont seg fault.
				if self.active_idx > self.chunk_size:
					self.data.resize((self.active_idx,1))
					self.data.attrs['blank_lines'] = True
				
			# save meta data.
			self.data.attrs['active_idx'] = self.active_idx
			self.data.attrs['sort_type'] = "none"
			
				
			# close file pointer.
			self.h5_out.close()		
			
		elif self.read == 1:
			self.h5_in.close()

		elif self.read == 2:
			self.h5_in.close()

	def load_complete(self):
		return self.data[:]
		
	def save_complete(self, a):
		if a.size == self.data.len():
			self.data[:] = a
		elif a.size > self.data.len():
			self.data.resize((a.size,1))
			self.data[:] = a
		else:
			self.data[0:a.size] = a
			
	def get_pointer(self):
		return self.data

	# save chunk to disk.
	def save_array(self, a):
		# increase size of h5 dataset.
		new_ptr = self.active_idx + a.size
		if new_ptr > self.max_idx:
			# grow h5 datasets.
			self.data.resize((new_ptr,1))
			
			# increase pointer.
			self.max_idx = new_ptr
			
		# save data to disk.
		#print "save_array"
		#print a.shape
		#print self.data.shape
		#print "end"
		self.data[self.active_idx:new_ptr] = a

		# update final pointer.
		self.active_idx = new_ptr

		
class HamAnnotate(object):
	
	# initializes h5 file.
	def __init__(self, ham_file, read=0):
		# decide how to open.
		if read == False:
			# Open h5 file for writing.
			self.h5_out = h5py.File(ham_file, 'w')
			self.read = False
		else:
			# open h5 file for reading.
			self.h5_ant_in = h5py.File(ham_file, 'r')
			self.read = True
	
	# gets pointers.
	def load_annotation(self):
		return self.h5_ant_in['ant_a'][:], self.h5_ant_in['ant_b'][:]
	
	# saves an annotation object.
	def save_annotation(self, a, b):
		# find chunk size.
		chunk_size = default_chunk_size
		if default_chunk_size > a.size:
			chunk_size = a.size
		if default_chunk_size > b.size:
			chunk_size = b.size
			
		
		# initialize dataset.
		self.data_a = self.h5_out.create_dataset('ant_a', (a.size,1), ham_ant_dt, compression='gzip', chunks=(chunk_size,1))
		self.data_b = self.h5_out.create_dataset('ant_b', (b.size,1), ham_ant_dt, compression='gzip', chunks=(chunk_size,1))

		# save dataset.
		self.data_a[:] = a
		self.data_b[:] = b
		
	# close annotation.
	def close(self):
		if self.read == 0:
			self.h5_out.close()
		else:
			self.h5_ant_in.close()
